import os
import cv2
import numpy as np
from tqdm import tqdm
import random

def add_floating_speckle_noise(image, speckle_prob=0.005, speckle_size=1):
    """
    添加漂浮物离散点噪声（随机黑色斑点）
    :param image: 单通道灰度图像
    :param speckle_prob: 每个像素点成为噪声点的概率
    :param speckle_size: 噪声点的扩散大小（1=单像素，2=2x2区域等）
    :return: 带噪声的图像
    """
    noisy = image.copy()
    h, w = image.shape
    
    # 生成随机噪声点位置
    noise_mask = np.random.rand(h, w) < speckle_prob
    
    # 将噪声点设为0（黑色）
    noisy[noise_mask] = 0
    
    # 可选：扩散噪声点（模拟稍大的漂浮物）
    if speckle_size > 1:
        kernel = np.ones((speckle_size, speckle_size), np.uint8)
        noise_mask_expanded = cv2.dilate(noise_mask.astype(np.uint8), kernel)
        noisy[noise_mask_expanded > 0] = 0
    
    return noisy

def process_images(
        input_dir, 
        output_dir,
        speckle_prob=0.005,
        speckle_size_lower=1,
        speckle_size_upper=2
    ):
    """
    批量处理图像并添加漂浮物噪声
    """
    # 噪声参数调整建议：
    # - speckle_prob: 0.001-0.01 (0.1%-1%的像素变为噪声点)
    # - speckle_size: 1-3 (模拟不同大小的漂浮物)
    speckle_size = random.randint(speckle_size_lower, speckle_size_upper)
    
    os.makedirs(output_dir, exist_ok=True)
    image_files = [f for f in os.listdir(input_dir) 
                  if f.lower().endswith(('.png'))]
    
    for filename in tqdm(image_files, desc="Adding floating speckle noise", unit="image"):
        try:
            img_path = os.path.join(input_dir, filename)
            img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
            
            if img is None:
                print(f"Warning: Could not read image {filename}, skipping...")
                continue
            
            # 添加漂浮物噪声
            noisy_img = add_floating_speckle_noise(img, speckle_prob, speckle_size)
            
            # 保存结果
            output_path = os.path.join(output_dir, filename)
            cv2.imwrite(output_path, noisy_img)
            
        except Exception as e:
            print(f"Error processing {filename}: {str(e)}")

if __name__ == "__main__":
    input_directory = "dataset/sample"  # 输入文件夹路径
    output_directory = "dataset/sample_with_noise"  # 输出文件夹路径

    process_images(input_directory, output_directory)